#include <ros/ros.h>
#include <ros/package.h>
#include <message_filters/subscriber.h>
#include <message_filters/time_synchronizer.h>
#include <message_filters/sync_policies/approximate_time.h>
#include <geometry_msgs/PoseStamped.h>
#include <sensor_msgs/Imu.h>

#include <vector>
#include <string>
#include <iostream>
#include <fstream>

#include <eigen3/Eigen/Dense>

using namespace std;

std::vector<Eigen::Matrix3d> v_v2m; // vicon to mav
std::vector<Eigen::Matrix3d> v_w2p; // world to px4

Eigen::Matrix3d m_w2v;
Eigen::Matrix3d m_m2p;

Eigen::Quaterniond last_q{1, 0, 0, 0};

int max_size = 100;
int max_iter = 100;
double min_q_diff = 20.0;
bool is_compare = true;

class rotationEstimator
{
public:
  rotationEstimator(){};
  ~rotationEstimator(){};

  bool setInput(const vector<Eigen::Matrix3d> &q_v2m,
                const vector<Eigen::Matrix3d> &q_w2p)
  {
    if (q_v2m.size() != q_w2p.size()) {
      cerr << "q_v2m.size() != q_w2p.size()" << endl;
      return false;
    }

    n_ = q_v2m.size();
    q_v2m_ = q_v2m;
    q_w2p_ = q_w2p;
    e_w2p_.clear();
    for (auto q : q_w2p) {
      e_w2p_.push_back(q2e(Eigen::Quaterniond(q)));
    }

    return true;
  }

  void estimate(int iter_times)
  {
    m_w2v_ = Eigen::Matrix3d::Identity();
    m_m2p_ = Eigen::Matrix3d::Identity();

    Eigen::MatrixXd A(3, n_), B(3, n_);

    for (size_t i = 0; i < n_; i++)
    {
      A.col(i) = Eigen::AngleAxisd(e_w2p_[i](0), Eigen::Vector3d::UnitX())
                     .toRotationMatrix()
                     .transpose() *
                 Eigen::AngleAxisd(e_w2p_[i](1), Eigen::Vector3d::UnitY())
                     .toRotationMatrix()
                     .transpose() *
                 Eigen::Vector3d::UnitZ();
    }

    for (size_t j = 0; j < iter_times; j++) {
      for (size_t i = 0; i < n_; i++) {
        B.col(i) = (m_w2v_ * q_v2m_[i]).transpose() * Eigen::Vector3d::UnitZ();
      }

      m_m2p_ = (B * A.transpose() * (A * A.transpose()).inverse());

      double x = 0.0;
      for (size_t i = 0; i < n_; i++) {
        Eigen::Vector3d v;
        v = q_v2m_[i] * m_m2p_ * A.col(i);
        x += atan2(v(1), v(2));
      }
      x = x / n_; // 这样比梯度下降鲁棒

      Eigen::Matrix3d _m_x =
          Eigen::AngleAxisd(x, Eigen::Vector3d::UnitX()).toRotationMatrix();
      double y = 0.0;
      for (size_t i = 0; i < n_; i++) {
        Eigen::Vector3d v;
        v = _m_x * q_v2m_[i] * m_m2p_ * A.col(i);
        y += -atan2(v(0), v(2));
      }
      y = y / n_;

      Eigen::Matrix3d _m_y =
          Eigen::AngleAxisd(y, Eigen::Vector3d::UnitY()).toRotationMatrix();
      m_w2v_ = _m_y * _m_x;
    }
  }

  static Eigen::Vector3d q2e(Eigen::Quaterniond q)
  {
    Eigen::Matrix3d R = q.toRotationMatrix();
    Eigen::Vector3d euler = R.eulerAngles(2, 1, 0);
    return {euler(2), euler(1), euler(0)};
  }

  void print()
  {
    cout << "m_w2v_: " << endl
         << q2e(Eigen::Quaterniond(m_w2v_)) << endl;
    cout << "m_m2p_: " << endl
         << q2e(Eigen::Quaterniond(m_m2p_)) << endl;
  }

  void getResults(Eigen::Matrix3d &m_w2v, Eigen::Matrix3d &m_m2p)
  {
    m_w2v = m_w2v_;
    m_m2p = m_m2p_;
  }

private:
  size_t n_{0};
  vector<Eigen::Matrix3d> q_v2m_;
  vector<Eigen::Matrix3d> q_w2p_;
  vector<Eigen::Vector3d> e_w2p_;

  Eigen::Matrix3d m_w2v_;
  Eigen::Matrix3d m_m2p_;
};

double getAngleDiff(const Eigen::Quaterniond &q1, const Eigen::Quaterniond &q2)
{
  Eigen::Quaterniond q_rel = q1 * q2.inverse();
  double angle = 2 * acos(abs(q_rel.w()));
  return angle * (180.0 / M_PI);
}

void callback(const geometry_msgs::PoseStamped::ConstPtr &pose,
              const sensor_msgs::Imu::ConstPtr &imu)
{
  Eigen::Quaterniond q_vicon(pose->pose.orientation.w,
                             pose->pose.orientation.x,
                             pose->pose.orientation.y,
                             pose->pose.orientation.z);
  Eigen::Quaterniond q_imu(imu->orientation.w,
                           imu->orientation.x,
                           imu->orientation.y,
                           imu->orientation.z);

  if (v_v2m.size() >= max_size) {
    Eigen::Matrix3d m_w2p = m_w2v * q_vicon.toRotationMatrix() * m_m2p;
    Eigen::Vector3d vicon_euler = rotationEstimator::q2e(Eigen::Quaterniond(m_w2p));
    Eigen::Vector3d px4_euler = rotationEstimator::q2e(q_imu);
    cout << "vicon: " << vicon_euler.transpose() *180.0/M_PI<< endl;
    cout << "px4:   " << px4_euler.transpose() *180.0/M_PI<< endl;
    return;
  }

  if (getAngleDiff(q_vicon, last_q) > min_q_diff) {
    v_v2m.push_back(q_vicon.toRotationMatrix());
    v_w2p.push_back(q_imu.toRotationMatrix());
    last_q = q_vicon;
  }
}

int main(int argc, char *argv[])
{
  ros::init(argc, argv, "calibrate");
  ros::NodeHandle nh;

  const string &node_name = ros::this_node::getName();
  nh.getParam(node_name + "/" + "max_size", max_size);
  nh.getParam(node_name + "/" + "max_iter", max_iter);
  nh.getParam(node_name + "/" + "min_q_diff", min_q_diff);
  nh.getParam(node_name + "/" + "is_compare", is_compare);

  cout << "[calibrate] max_size:   " << max_size << endl;
  cout << "[calibrate] max_iter:   " << max_iter << endl;
  cout << "[calibrate] min_q_diff: " << min_q_diff << endl;
  cout << "[calibrate] is_compare: " << is_compare << endl;

  message_filters::Subscriber<geometry_msgs::PoseStamped> sub1(nh, "/pose", 10);
  message_filters::Subscriber<sensor_msgs::Imu> sub2(nh, "/mavros/imu/data", 10);
  typedef message_filters::sync_policies::ApproximateTime<
      geometry_msgs::PoseStamped, sensor_msgs::Imu>
      MySyncPolicy;
  message_filters::Synchronizer<MySyncPolicy> sync(MySyncPolicy(10), sub1, sub2);
  sync.registerCallback(boost::bind(&callback, _1, _2));

  while (ros::ok()) {
    ros::spinOnce();

    if (v_v2m.size() >= max_size) break;

    ROS_INFO_STREAM_THROTTLE(
        0.5, "collecting data..." << v_v2m.size() << "/" << max_size);

    ros::Rate(25).sleep();
  }

  ROS_INFO("data collected, start calibration...");

  rotationEstimator estimator;
  estimator.setInput(v_v2m, v_w2p);
  estimator.estimate(100);
  estimator.print();
  estimator.getResults(m_w2v, m_m2p);

  // record the quaternion of m_w2v and m_m2p into csv in ../cfg/calibration.csv
  string backage_name = "viconros";
  string file_path = ros::package::getPath(backage_name) + "/cfg/calibration.csv";
  ofstream file(file_path);

  // record to file in form of: w, x, y, z
  Eigen::Quaterniond q_w2v(m_w2v);
  Eigen::Quaterniond q_m2p(m_m2p);
  file << "q_w2v" << endl;
  file << q_w2v.w() << ", " << q_w2v.x() << ", "
       << q_w2v.y() << ", " << q_w2v.z() << endl;
  file << "q_m2p" << endl;
  file << q_m2p.w() << ", " << q_m2p.x() << ", "
       << q_m2p.y() << ", " << q_m2p.z() << endl;

  ROS_INFO("calibration done");

  if (is_compare) {
    ros::spin();
  }

  return 0;
}
